package org.dromara.common.encrypt.utils;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * AES秘钥生成
 * @author Stars
 */
public class AESKeyGenerator {

    /**
     * 生成AES密钥
     * @param keySize 密钥长度，支持128、192或256位
     * @return 生成的密钥对象
     * @throws Exception 如果生成过程中出错
     */
    public static SecretKey generateKey(int keySize) throws Exception {
        // 检查密钥长度是否合法
        if (keySize != 128 && keySize != 192 && keySize != 256) {
            throw new IllegalArgumentException("AES密钥长度必须是128、192或256位");
        }

        // 获取AES密钥生成器
        KeyGenerator keyGen = KeyGenerator.getInstance("AES");

        // 使用安全随机数初始化生成器
        SecureRandom secureRandom = SecureRandom.getInstanceStrong();
        keyGen.init(keySize, secureRandom);

        // 生成密钥
        return keyGen.generateKey();
    }

    /**
     * 将SecretKey转换为Base64编码的字符串
     * @param key 密钥对象
     * @return Base64编码的字符串
     */
    public static String keyToBase64(SecretKey key) {
        return Base64.getEncoder().encodeToString(key.getEncoded());
    }

    /**
     * 从Base64编码的字符串恢复SecretKey
     * @param base64Key Base64编码的密钥字符串
     * @return 恢复的密钥对象
     */
    public static SecretKey keyFromBase64(String base64Key) {
        byte[] keyBytes = Base64.getDecoder().decode(base64Key);
        return new javax.crypto.spec.SecretKeySpec(keyBytes, "AES");
    }

    public static void main(String[] args) {
        try {
            // 生成256位AES密钥
            SecretKey secretKey = generateKey(128);

            // 转换为Base64字符串以便存储或传输
            String base64Key = keyToBase64(secretKey);
            System.out.println("生成的AES密钥(Base64): " + base64Key);

            // 从Base64字符串恢复密钥
            SecretKey restoredKey = keyFromBase64(base64Key);
            System.out.println("恢复的密钥与原密钥是否相同: " +
                java.util.Arrays.equals(secretKey.getEncoded(), restoredKey.getEncoded()));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
